
from sqlalchemy import create_engine,table,Column,Integer,String,MetaData
from sqlalchemy.orm import sessionmaker,declarative_base

from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage,HumanMessage
from langgraph.graph import StateGraph

from pydantic import BaseModel,Field
from typing import Optional,Union,TypedDict,Annotated

import operator


Base = declarative_base()

class UserInfo(Base):
    __tablename__ = "users"
    id = Column(Integer,primary_key=True)
    name = Column(String(50))
    age = Column(Integer)
    email = Column(String(50))
    phone = Column(String(20))


# 数据库连接 URI，这里要替换成自己的Mysql 连接信息，以下是各个字段的对应解释：
# root：MySQL 数据库的用户名。
# 123456：MySQL 数据库的密码。
# localhost：MySQL 服务器的 IP 地址。
# langgraph：要连接的数据库的名称。
# charset=utf8mb4：设置数据库的字符集为 utf8mb4，支持更广泛的 Unicode 字符
DATABASE_URI = "mysql+pymysql://root:1qaz2wsx@localhost:3306/langgraph?charset=utf8mb4"

engine = create_engine(DATABASE_URI,echo=True)

# 如果表不存在，则创建表
Base.metadata.create_all(engine)
# 创建会话
Session = sessionmaker(bind=engine)
session = Session()

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-4o",api_key = api_key,base_url=api_base,temperature=0)

class User(BaseModel):
    name:str=Field(description="用户名称")
    age:Optional[int]=Field(description="用户年龄")
    email:str=Field(description="用户邮箱")
    phone:Optional[str]=Field(description="用户手机号")

class ConversationalResponse(BaseModel):
    response: str = Field(description="A conversational response to the user's query")


class FinalResponse(BaseModel):
    final_output: Union[User,ConversationalResponse]

def chat_with_model(state):
    print(state)
    messages = state["messages"]
    structured_llm = llm.with_structured_output(FinalResponse)
    resp = structured_llm.invoke(messages)
    return {"messages":[resp]}

def final_answer(state):
    print(state)
    messages = state["messages"][-1]
    resp = messages.final_output.response
    return {"messages":[resp]}

def insert_db(state):
    # 确保为每次操作创建新的会话
    session = Session()
    try:
        result = state["messages"][-1]
        output = result.final_output
        # 创建用户实例
        user = UserInfo(name=output.name,age=output.age,email=output.email,phone=output.phone)
        # 添加到会话
        session.add(user)
        # 提交事务
        session.commit()
        return {"messages":[f"数据已经成功存储到数据库"]}
    except Exception as e:
        session.rollback()
        return {"messages":[f"数据存储失败，错误原因:{e}"]}
    finally:
        session.close()


class AgentState(TypedDict):
    messages:Annotated[list[AnyMessage],operator.add]

def generate_branch(state:AgentState):
    result = state["messages"][-1]
    output = result.final_output
    
    if isinstance(output,User):
        return True
    else:
        return False
    
builder = StateGraph(AgentState)

builder.add_node("chat_with_model",chat_with_model)
builder.add_node("final_answer",final_answer)
builder.add_node("insert_db",insert_db)

# 设置图的启动节点
builder.set_entry_point("chat_with_model")

# 设置条件边
builder.add_conditional_edges("chat_with_model",generate_branch,{True:"insert_db",False:"final_answer"})

# 设置终止节点
builder.set_finish_point("final_answer")
builder.set_finish_point("insert_db")

graph = builder.compile()

query="我叫奥特曼，今年38岁，邮箱地址是aoteman#qq.com，电话是1211111111"
query1 = "介绍下自己"

input_message = {"messages":[HumanMessage(content=query1)]}
resp = graph.invoke(input_message)
print(resp)

